import zmq
import ioloop
import logging
import os


zmq_transport_INPROC = 0
zmq_transport_IPC = 1
zmq_transport_TCP = 2
zmq_transport_PGM = 3


class ZNodeError(Exception):
    """docstring for ZNodeException"""
    def __init__(self, arg):
        super(ZNodeError, self).__init__(arg)
        self.arg = arg
        self.str = "ZNodeError"

    def __str__(self):
        """docstring for __str__"""
        return "%s : %s" % (self.str, str(self.arg))
    #__str__()
#ZNodeException


class ZNodeStateError(ZNodeError):
    """docstring for ZNodeStateError"""
    def __init__(self, arg):
        super(ZNodeStateError, self).__init__(arg)
        self.str = "Unknown bad state or event out of expected order"
#ZNodeStateError


class SocketTypeUnkownError(ZNodeError):
    def __init__(self, arg):
        super(SocketTypeUnkownError, self).__init__(arg)
        self.str = "Unknown ZeroMQ socket type"
#SocketTypeUnkownError


class TransportUnkownError(ZNodeError):
    def __init__(self, arg):
        super(TransportUnkownError, self).__init__(arg)
        self.str = "Unkown ZeroMQ transport"
#TransportUnkownError


class SendNotSupportedError(ZNodeError):
    def __init__(self, arg):
        super(SendNotSupportedError, self).__init__(arg)
        self.str = "ZNode type does not support sends"
#SendNotSupportedError


class RecvNotSupportedError(ZNodeError):
    def __init__(self, arg):
        super(RecvNotSupportedError, self).__init__(arg)
        self.str = "ZNode type does not support receives"
#RecvNotSupportedError


class ZNodeUnknownEventError(ZNodeError):
    def __init__(self, arg):
        super(ZNodeUnknownEventError, self).__init__(arg)
        self.str = "Received an unkown event"
#UZNodenknownEventError


class ZNodePollError(ZNodeError):
    def __init__(self, arg):
        super(ZNodePollError, self).__init__(arg)
        self.str = "ZNode Poll Error"
#ZNodePollError


class ZNode(object):
    """docstring for ZNode"""

    POLLIN = zmq.core.poll.POLLIN
    POLLOUT = zmq.core.poll.POLLOUT
    POLLERR = zmq.core.poll.POLLERR

    type_map = {
        zmq.REQ: 'REQ',
        zmq.REP: 'REP',
        zmq.PUB: 'PUB',
        zmq.SUB: 'SUB',
        zmq.PUSH: 'PUSH',
        zmq.PULL: 'PULL',
        zmq.ROUTER: 'ROUTER',
        zmq.DEALER: 'DEALER',
    }

    trans_map = {
        zmq_transport_INPROC: 'INPROC',
        zmq_transport_IPC: 'IPC',
        zmq_transport_TCP: 'TCP',
        zmq_transport_PGM: 'PGM'
    }

    def __init__(self,
              app=None,
              transport=zmq_transport_TCP,
              endpoint='*',
              port=5555,
              context=None):

        self.socket = None
        self.event = None
        self._send_kargs = {}
        self._send_resp = None
        self._on_recv_func = None
        self._on_recv_func_persist = False
        self._on_send_func = None
        self._on_send_func_persist = False
        self._outbound = None
        self._state = 0

        self.ioloop = ioloop.IOLoop.instance()
        self.app = None
        self.type = zmq.REP
        self.transport = transport
        self.endpoint = endpoint
        self.port = port
        self.context = context or zmq.Context.instance()

        # we'll try to send this many times in a row before dropping the
        # send data
        self.max_sends = 10
        self.max_sends_tries = 0

        self.device = ''
        self.environ = dict(os.environ.items())

        if self.transport == zmq_transport_INPROC:
            self.device = 'inproc://%s' % (self.endpoint,)
        elif self.transport == zmq_transport_IPC:
            self.device = 'ipc://%s' % (self.endpoint,)
        elif self.transport == zmq_transport_TCP:
            self.device = 'tcp://%s:%s' % (self.endpoint, self.port)
        elif self.transport == zmq_transport_PGM:
            self.device = 'pgm://%s:%s' % (self.endpoint, self.port)
        else:
            raise TransportUnkownError(self.transport)

        self.environ['DMQ_IOLOOP'] = self.ioloop
        self.environ['ZMQ_ENDPIONT'] = self.endpoint
        self.environ['ZMQ_TRANSPORT'] = self.transport
        self.environ['ZMQ_PORT'] = self.port
        self.environ['ZMQ_CONTEXT'] = self.context
        self.environ['ZMQ_TYPE'] = self.type

        self.set_app(app)

        logging.debug("device:%s" % (self.device,))
    #__init__()

    def set_app(self, app):
        """docstring for set_app"""
        logging.debug("ZNode.set_app(self : %s, app : %s)" % (self, app))
        self.app = app
        if self.app and self.socket and not self.socket.closed:
            self.update(ZNode.POLLIN)

        return self
    #set_app()

    def bind(self):
        logging.debug("ZNode.bind() type:%s" % (self.type,))
        self.socket = self.context.socket(self.type)
        self.socket.bind(self.device)

        return self
    #bind()

    def connect(self):
        logging.debug("ZNode.connect, device:%s, type:%s"
            % (self.device, ZNode.type_map[self.type],))
        self.socket = self.context.socket(self.type)
        self.socket.connect(self.device)

        return self
    #connect()

    def close(self):
        """docstring for close"""
        self.ioloop.remove_node(self)
        self.socket.close()
        self.socket = None
    #close()

    def add(self):
        """docstring for add"""
        logging.debug("ZNode.add() %s " % (self,))
        if not self.socket:
            raise ZNodeError("ZNode.add() socket : %s, non existent"
                            % (self.socket,))

        # See if we have any client callbacks set and if so
        # set up the proper events.
        events = 0
        if self.app or self._on_recv_func:
            logging.debug("ZNode.add() setting POLLIN")
            events |= ZNode.POLLIN
        if self._on_send_func:
            logging.debug("ZNode.add() setting POLLOUT")
            events |= ZNode.POLLOUT

        self._state |= events

        self.ioloop.add_node(self, self._state)
        return self
        #return self.update(events)
    #add()

    def _recv_handler(self, event):
        """docstring for _recv_hanlder"""
        raise RecvNotSupportedError(
                "_recv_handler has not been implemented by this class")
    #_recv_hanlder()

    def _send_handler(self, event):
        """docstring for _send_handler"""
        raise SendNotSupportedError(
                "_send_handler has not been implemented by this class")
    #_send_handler()

    def _err_handler(self, event):
        """docstring for _err_handler"""
        logging.debug("_err_handler")
        # Need to figure out how to get the error information from
        # ZMQ and the IOLoop so we can construct a better exception message
        raise ZNodePollError()
    #_err_handler()

    def poll_handler(self, event):
        """docstring for poll_handler"""

        logging.debug("poll_handler(event:%2.2X)" % (event,))
        if event & ZNode.POLLIN:
            self._recv_handler(event)

        if event & ZNode.POLLOUT:
            self._send_handler(event)

        if event & ZNode.POLLERR:
            self._err_handler(event)
    #poll_handler()

    def send(self, data, flags=0, copy=True, track=False):
        raise SendNotSupportedError()
    #send()

    def on_send(self, func):
        raise SendNotSupportedError()
    #on_send()

    def on_recv(self, func):
        raise RecvNotSupportedError()
    #on_recv()

    def update(self, events, rem=False):
        logging.debug(("ZNode.update(self._state:%2.2X,"
                    " events:%2.2X, rem=%s)")
                    % (self._state, events, rem))

        prev_state = self._state
        if not rem and not ((self._state & events) == events):
            # Only update if the new event is not already
            # registered
            logging.debug("ZNode.update adding events %2.2X)" % (events,))
            self._state |= events
        elif rem and (self._state & events):
            # If we have any of the events that need to be removed
            # we act
            logging.debug("ZNode.update removing events %2.2X)" % (events,))
            self._state &= ~events

        if prev_state != self._state:
            logging.debug("ZNode.update updating to self._state:%2.2X)"
                        % (self._state,))
            self.ioloop.update_node(self, self._state)

        return self
    #update()

    def __str__(self):
        cname = self.__class__.__name__
        str = "%s, socket:%s" % (cname, self.socket)
        str += ("\n_state : %2.2X, type : %s, device : %s"
                % (self._state, ZNode.type_map[self.type], self.device))

        if self.app:
            str += "\napp : %s" % (self.app,)
        if self._on_recv_func:
            str += "\n_on_recv_func : %s" % (self._on_recv_func,)
        if self._on_send_func:
            str += "\n_on_send_func : %s" % (self._on_send_func,)

        return str
    #__str__()

    def __repr__(self):
        """docstring for __repr__"""
        return self.__str__()
    #__repr__()
#ZNode


class ZNodeSender(object):

    def __init__(self, *args, **kwargs):
        super(ZNodeSender, self).__init__(*args, **kwargs)
    #__init__()

    def add(self):
        logging.debug("add() connect, %s" % self)
        self.connect()
        return ZNode.add(self)
    #add()

    def _send_handler(self, event):
        """docstring for _send_handler"""
        logging.debug("ZNodeSender._send_handler(event:%2.2X)" % (event,))

        if self._on_send_func:
            # The _on_send_func can be set during a call to _on_send_func
            sfunc = self._on_send_func
            if not self._on_send_func_persist:
                self._on_send_func = None
            self._outbound = sfunc(self.environ)

        if self._outbound:
            logging.debug(("ZNodeSender._send_handler() sending"
                        "_outbound:%s, _send_kargs:%s")
                        % (self._outbound, self._send_kargs))
            try:
                self.socket.send(self._outbound, **self._send_kargs)
            except Exception as e:
                if self.max_sends_tries < self.max_sends:
                    self.max_sends_tries += 1
                    # we go ahead and pass the Exeption up,
                    # it will be handled and reported
                    # by the ioloop
                    raise

                # This is fairly bad, it's probably because
                # we're trying to send something
                # that is not sendable, no buffer interface,
                # or there is a zmq problem.
                # If we don't discard the data, we'll keep
                # trying to send it and get stuck
                # in a send error loop.
                logging.error(("ZNodeSender._send_handler() Exception while "
                            "sending, data is being dropped."
                            "data : %s\n, Exception : %s\n")
                            % (self._outbound, e), exc_info=True)

            self.max_sends_tries = 0
            self._outbound = None

        # we entered here with POLLOUT set, don't disable it
        # unless _on_send_func is not set.
        if not self._on_send_func:
            self.update(ZNode.POLLOUT, rem=True)

        logging.debug(
                "ZNodeSender._send_handler() done, _outbound:%s, _state:%s"
                % (self._outbound, self._state))
    #_send_handler()

    def on_send(self, func, persist=False):
        logging.debug("ZNodeSender.on_send() func:%s, _state:%2.2X"
                % (func, self._state))

        self._on_send_func = func
        self._on_send_func_persist = persist

        if self._on_send_func:
            self.update(ZNode.POLLOUT)
        else:
            self.update(ZNode.POLLOUT, rem=True)

        return self
    #on_send()

    def send(self, data, flags=0, copy=True, track=False):
        """
        send() can also be used to cancel data and/or set up args
        other than data that can be used on later sends.
        """
        logging.debug("ZnodeSender.send(data:%s, flags:%s)" % (data, flags))

        self._send_kargs = {'flags': flags,
                          'copy': copy, 'track': track}
        self._outbound = data
        if self._outbound:
            self.update(ZNode.POLLOUT)
        elif ZNode.POLLOUT == (ZNode.POLLOUT & self._state):
            # avoid the update call if we can here
            self.update(ZNode.POLLOUT, rem=True)

        return self
    #send()

#ZNodeSender


class ZNodeRecvr(object):

    def add(self):
        logging.debug("add() bind, %s" % self)
        self.bind()
        return ZNode.add(self)
    #add()

    def on_recv(self, func=None, persist=False):
        """
        Add a recieve handler into the ioloop. When the socket is ready to
        recv func will be called with the data that was recved on the socket.

        calling on_recv(), with no params, will cancel any set on_recv function
        and set persistent to False.
        """

        logging.debug("on_recv(func:%s, persist:%s) %s"
                    % (func, persist, self))
        logging.debug("on_recv() app:%s, _on_recv_func:%s"
                    % (self.app, self._on_recv_func))
        self._on_recv_func_persist = persist
        self._on_recv_func = func

        if self._on_recv_func or self.app:
            self.update(ZNode.POLLIN)
        else:
            self.update(ZNode.POLLIN, rem=True)

        return self
    #on_recv()

    def _recv_handler(self, event):
        """docstring for _recv_hanlder"""

        #Need to update this handler to use the single/multi recv
        #pattern so it can recv both types of messages correctly.
        #Even with poll, if a multipart msg is recvd, we can get
        #it all in one call of _recv_handler.

        logging.debug("_recv_handler(event:%2.2X) %s" % (event, self))
        # process the request. If the
        # app handler/on_recv returns data, it's assumed
        # to be an outbound request

        inbound = self.socket.recv()
        logging.debug("_recv_handler(event:%2.2X), %s, recv'd : %s"
                    % (event, self, inbound))

        if self._on_recv_func:
            # The _on_recv_func can be set in the _on_recv_func
            rfunc = self._on_recv_func
            if not self._on_recv_func_persist:
                self._on_recv_func = None
            self._outbound = rfunc(inbound, self.environ)
        elif self.app:
            self._outbound = self.app(inbound, self.environ)
        else:
            self.update(ZNode.POLLIN, rem=True)
            raise ZNodeStateError(
                        ("_recv_handler : Processing POLLIN, but there "
                        "is no app() handler or on_recv() registered"))

        # during the clients recv handler the client can
        # remove app or _on_recv_func if they wish.
        if not (self.app or self._on_recv_func):
            logging.debug(
                    "_recv_handler() no client handler, removing POLLIN %s "
                    % (self,))
            self.update(ZNode.POLLIN, rem=True)

    #_recv_hanlder()
#ZNodeRecvr


class ZNodeREQ(ZNodeSender, ZNodeRecvr, ZNode):

    """docstring for ZNodeREQ"""
    def __init__(self, *args, **kwargs):
        super(ZNodeREQ, self).__init__(*args, **kwargs)
        self.type = zmq.REQ
    #__init__()

    def add(self):
        self.connect()
        return ZNode.add(self)
    #add()

    def _recv_handler(self, event):
        """docstring for _recv_hanlder"""
        super(ZNodeREQ, self)._recv_handler(event)
        if self._outbound:
            self.update(ZNode.POLLOUT)
    #_recv_hanlder()
#ZNodeREQ


class ZNodeREP(ZNodeRecvr, ZNodeSender, ZNode):
    """docstring for ZNodeREP"""
    def __init__(self, *args, **kwargs):
        super(ZNodeREP, self).__init__(*args, **kwargs)
        self.type = zmq.REP
    #__init__()

    def add(self):
        self.bind()
        return ZNode.add(self)
    #add()

    def _recv_handler(self, event):
        """docstring for _recv_hanlder"""
        super(ZNodeREP, self)._recv_handler(event)
        if self._outbound:
            self.update(ZNode.POLLOUT)
    #_recv_hanlder()
#ZNodeREP

#####################################################################
##### No sure what do with the dealer and router yet. They don't
##### actually generate poll events as far as I can tell. So the
##### ZNode associated with them may just be used for controll and
##### monitoring.
#####################################################################


class ZNodeDEALER(ZNodeRecvr, ZNode):
    """docstring for ZNodeDEALER"""
    def __init__(self, *args, **kwargs):
        super(ZNodeDEALER, self).__init__(*args, **kwargs)
        self.type = zmq.DEALER
    #__init__()
#ZNodeDEALER


class ZNodeROUTER(ZNodeRecvr, ZNode):
    """docstring for ZNodeROUTER"""
    def __init__(self, *args, **kwargs):
        super(ZNodeROUTER, self).__init__(*args, **kwargs)
        self.type = zmq.ROUTER
    #__init__()
#ZNodeROUTER
#####################################################################


class ZNodePUB(ZNodeSender, ZNode):
    """docstring for ZNodePUB"""
    def __init__(self, *args, **kwargs):
        super(ZNodePUB, self).__init__(*args, **kwargs)
        self.type = zmq.PUB
    #__init__()

    def add(self):
        self.bind()
        return ZNode.add(self)
    #add()

#ZNodePUB


class ZNodeSUB(ZNodeRecvr, ZNode):
    """docstring for ZNodeSUB"""
    def __init__(self, *args, **kwargs):
        super(ZNodeSUB, self).__init__(*args, **kwargs)
        self.type = zmq.SUB
    #__init__()

    def add(self, filter=""):
        self.connect()
        self.subscribe(filter)
        return ZNode.add(self)
    #add()

    def subscribe(self, filter=""):
        logging.debug("ZNodeSUB.subscribe(filter:'%s'" % (filter,))
        self.socket.setsockopt(zmq.SUBSCRIBE, filter)
        return self
    #subscribe()

    def unsubscribe(self, filter=""):
        """docstring for unsubscribe"""
        self.socket.setsockopt(zmq.UNSUBSCRIBE, filter)
        return self
    #unsubscribe()
#ZNodeSUB


class ZNodePUSH(ZNodeSender, ZNode):
    """docstring for ZNodePUSH"""
    def __init__(self, *args, **kwargs):
        super(ZNodePUSH, self).__init__(*args, **kwargs)
        self.type = zmq.PUSH
    #__init__()
#ZNodePUSH


class ZNodePULL(ZNodeRecvr, ZNode):
    """docstring for ZNodePULL"""
    def __init__(self, *args, **kwargs):
        super(ZNodePULL, self).__init__(*args, **kwargs)
        self.type = zmq.PULL
    #__init__()
#ZNodePULL


def znode(app=None, type=zmq.REP,
            transport=zmq_transport_TCP,
            endpoint='*',
            port=5555,
            context=None,
            auto=True):
    """docstring for znode"""

    # a hack here. If the transport is specified as a string
    # we need to get the index valud for it.
    for k, v in ZNode.trans_map.items():
        try:
            if transport.upper() == v:
                transport = k
                break
        except AttributeError:
            continue

    logging.debug(("znode(type:%s, transport:%s, endpoint:%s, "
                    "port:%s, context:%s, auto:%s)")
                    % (ZNode.type_map[type], ZNode.trans_map[transport],
                        endpoint, port, context, auto))

    node = None
    if type == zmq.REP:
        node = ZNodeREP(app, transport, endpoint, port, context)
    elif type == zmq.REQ:
        node = ZNodeREQ(app, transport, endpoint, port, context)
    elif type == zmq.DEALER:
        node = ZNodeDEALER(app, transport, endpoint, port, context)
    elif type == zmq.ROUTER:
        node = ZNodeROUTER(app, transport, endpoint, port, context)
    elif type == zmq.PUB:
        node = ZNodePUB(None, transport, endpoint, port, context)
    elif type == zmq.SUB:
        node = ZNodeSUB(app, transport, endpoint, port, context)
    elif type == zmq.PUSH:
        node = ZNodePUSH(None, transport, endpoint, port, context)
    elif type == zmq.PULL:
        node = ZNodePULL(app, transport, endpoint, port, context)
    else:
        raise SocketTypeUnkownError()

    if auto:
        node.add()

    return node
#znode()
